import tensorflow as tf
from keras_preprocessing.image import ImageDataGenerator
from PIL import Image
import cv2 as cv
import os


def png_to_jpg(filepath):
    """
    将png图片转换为jpg图片
    :param filepath: 图片路径
    :return: 转换之后的图片路径
    """
    img = cv.imread(filepath, 0)
    w, h = img.shape[::-1]
    infile = filepath
    outfile = os.path.splitext(infile)[0] + ".jpg"
    img = Image.open(infile)
    img = img.resize((int(w / 2), int(h / 2)), Image.ANTIALIAS)
    try:
        if len(img.split()) == 4:
            # prevent IOError: cannot write mode RGBA as BMP
            r, g, b, a = img.split()
            img = Image.merge("RGB", (r, g, b))
            img.convert('RGB').save(outfile, quality=100)
            os.remove(filepath)
        else:
            img.convert('RGB').save(outfile, quality=100)
            os.remove(filepath)
        return outfile
    except Exception as e:
        print("PNG转换JPG出现错误！", e)


def augment_image(x):
    """
    随机改变图像的色调和饱和度，以模拟不同的照明条件
    :param x: 图片矩阵
    :return: 新图片矩阵
    """
    x = tf.image.random_saturation(x, 0.9, 1.2)  # 饱和度
    x = tf.image.random_hue(x, 0.02)  # 色调
    return x


def convert_to_hsv_and_grayscale(x):
    """
    创建一个自定义图层，用于将原始图像从RGB转换为HSV和灰度，并连接结果，最终变为尺寸为(宽度,高度,4)
    :param x: 图片矩阵
    :return: 新图片矩阵
    """
    hsv = tf.image.rgb_to_hsv(x)
    gray = tf.image.rgb_to_grayscale(x)
    rez = tf.concat([hsv, gray], axis=-1)
    return rez


def build_data_generators(train_folder: str, test_folder: str, labels: [] = None, image_size=(100, 100), batch_size=50,
                          validation_split=0.1):
    """
    数据读取与生成器；
    给出训练和测试文件夹的路径以及验证和测试的比例，这个方法会创建三个生成器:
     - 训练生成器使用训练集中的(100-验证_百分比)的图像
    它应用随机的水平和垂直翻转来增加数据，并随机地生成批次
    - 验证生成器使用训练集中剩余的验证_百分比的图像
    不生成随机批次，因为模型没有在这些数据上进行训练
    使用验证数据监控准确性和损失，以便在模型达到局部最优时更新学习率
    - 测试生成器使用的是测试集，没有任何形式的增强。
    一旦训练过程结束，准确率和损失的最终值就会在这组数据上计算出来
    :param train_folder: 测试文件夹名称
    :param test_folder: 训练文件夹名称
    :param labels: 类别名称列表
    :param image_size: 图片规格
    :param batch_size: 训练批大小
    :param validation_split: 验证集划分比例
    :return: 训练数据集：train_gen, 验证数据集：valid_gen, 测试数据集：test_gen
    """
    train_data_generator = ImageDataGenerator(
        width_shift_range=0.0,
        height_shift_range=0.0,
        zoom_range=0.0,
        validation_split=validation_split,
        horizontal_flip=True,
        vertical_flip=True,  # 随机翻转图像
        preprocessing_function=augment_image)  # 扩增处理目前只在训练集上进行（也可选择在验证集上进行）
    test_data_generator = ImageDataGenerator()
    train_gen = train_data_generator.flow_from_directory(train_folder, target_size=image_size, class_mode='sparse',
                                                         batch_size=batch_size, shuffle=True, subset='training',
                                                         classes=labels)
    valid_gen = train_data_generator.flow_from_directory(train_folder, target_size=image_size, class_mode='sparse',
                                                         batch_size=batch_size, shuffle=False, subset='validation',
                                                         classes=labels)
    test_gen = test_data_generator.flow_from_directory(test_folder, target_size=image_size, class_mode='sparse',
                                                       batch_size=batch_size, shuffle=False, subset=None,
                                                       classes=labels)
    return train_gen, valid_gen, test_gen
